#include <algorithm>
#include <cmath>
#include <cstdlib>
// #include <functional>
// #include <iostream>
// #include <numeric>
#include <stdexcept>
#include <string>

#include "DoseStatCuda.h"
#include "MacrosCuda.h"
#include "driver_types.h"

// CUDA相关头文件
#include <cuda_runtime.h>
#include <device_launch_parameters.h>

#include "CudaFloats.cuh"
#include "DoseStatCuda.h"
#include "MyPrtFns.h"

/// 控制是否打印
#if 0
#    define MYDBG printf
#else
#    define MYDBG(...)
#endif

constexpr int BLOCK_SIZE = 256;

// 为计算能力低于6.0的设备添加double类型的atomicAdd支持
#if __CUDA_ARCH__ < 600
__device__ double atomicAddLow(double* address, double val)
{
    unsigned long long int* address_as_ull = (unsigned long long int*)address;
    unsigned long long int old = *address_as_ull, assumed;
    do
    {
        assumed = old;
        old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));
        // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
    } while (assumed != old);
    return __longlong_as_double(old);
}
#endif

//xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
namespace dosestatcuda
{
    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  设备函数: 分箱统计体积
    __global__ void
    binVolumeKernel(const double* doses, const double* volumes, size_t num_pstv, double minDose, double binWidth, int numBins, double* binVolumes)
    {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx >= num_pstv) return;

        double dose = doses[idx];
        double volume = volumes[idx];
        int binIndex;

        if (dose <= minDose)
        {
            binIndex = 1;  // 哨兵0
        }
        else if (dose >= minDose + (numBins - 1) * binWidth)
        {
            // 哨兵n+1, 这里我们不统计这部分体积
            return;
        }
        else
        {
            // 落点取bin的右边界, 对应高的dose
            binIndex = static_cast<int>((dose - minDose) / binWidth + 0.5) + 1;
        }
#if __CUDA_ARCH__ < 600
        atomicAddLow(&binVolumes[binIndex], volume);
#else
        atomicAdd(&binVolumes[binIndex], volume);
#endif
    }

    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 设备函数: 计算累积相对体积
    __global__ void computeAccumulatedVolumesKernel(double* accumulatedRelativeVolumes, size_t numBins)
    {
        // 使用单一线程计算累积体积，因为这是一个顺序依赖的操作
        if (threadIdx.x == 0 && blockIdx.x == 0)
        {
            double totalVol = 0.0;
            for (int i = numBins; i > 0; --i)
            {
                double vol = accumulatedRelativeVolumes[i];
                accumulatedRelativeVolumes[i] = totalVol + 0.5 * vol;
                totalVol += vol;
            }

            // 设置哨兵值
            accumulatedRelativeVolumes[0] = totalVol;
            accumulatedRelativeVolumes[numBins + 1] = 0.0;

            // 归一化
            if (totalVol > 0.0)
            {
                for (size_t i = 0; i < numBins + 2; ++i)
                {
                    accumulatedRelativeVolumes[i] /= totalVol;
                }
            }
        }
    }

    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  设备函数: 计算加权剂量（用于平均剂量计算）
    __global__ void ComputeWeightedDoseKernel(const double* voxelVolume,
                                              const double* voxelDose,
                                              size_t size,
                                              double* totalVolume,
                                              double* weightedDose)
    {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx >= size) return;

        double vol = voxelVolume[idx];
        double dose = voxelDose[idx];

        atomicAdd(totalVolume, vol);
        atomicAdd(weightedDose, dose * vol);
    }

    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  计算DVH, dose-volume histogram (X-Y) - CUDA版本
    DoseStatRet Accumulate(const std::vector<double>& voxelVolume, const std::vector<double>& voxelDose, bool doExactAccumulation)
    {
        MYDBG("doExactAccumulation: %d\n", doExactAccumulation);
        DoseStatRet result;

        // 预分配可能的内存以减少动态扩容; dose->volume
        const size_t numVoxels = voxelVolume.size();
        std::vector<std::pair<double, double>> doseVolumePairs;
        doseVolumePairs.reserve(numVoxels);

        /// 过滤体积为0的体素，同时直接创建 剂量->体积 对
        for (size_t i = 0; i < numVoxels; ++i)
        {
            if (voxelVolume[i] > 0)
            {
                doseVolumePairs.emplace_back(std::max(0.0, voxelDose[i]), voxelVolume[i]);
            }
        }
        doseVolumePairs.shrink_to_fit();
        // number of positive voxels
        const size_t num_pstv = doseVolumePairs.size();

        ///==== 空数据的快速处理
        if (num_pstv == 0)
        {
            /// ticks 最少为 2 => 一个 bin
            /// volume; 首哨兵为总体积, 末哨兵为 0.0
            /// dose; 首哨兵为 min dose, 末哨兵为 max dose
            result.accumulatedRelativeVolumes = {1.0, 0.0};
            result.binnedDoseValues = {0.0, 0.0};

            return result;
        }

        /// number of blocks
        const int nBlocks = MYCEIL(num_pstv, BLOCK_SIZE);
        //MYINFO("[my] num_pstv -> {}, blocks -> {}", num_pstv, nBlocks);

        // 根据数据量和doExactAccumulation参数选择合适的算法
        constexpr size_t EXACT_THRESHOLD = 5000;
        constexpr int MAX_BINS = 1001;
        constexpr int MIN_BINS = 2;
        constexpr double doseMinWidth = 1e-6;  // dose bin 最小分辨率
        const bool useExactMethod = doExactAccumulation || (num_pstv < EXACT_THRESHOLD);

        do
        {
            ///============================= 精确计算：
            /// 直接排序 {剂量->体积}, 先按照剂量排序，然后按照体积排序
            if (useExactMethod)
            {
                // 避免额外的索引数组和多次拷贝
                std::sort(doseVolumePairs.begin(), doseVolumePairs.end());

                // 准备结果数组
                result.binnedDoseValues.resize(num_pstv + 2, 0.0);
                result.accumulatedRelativeVolumes.resize(num_pstv + 2, 0.0);

                // 填充排序后的结果
                for (size_t i = 0; i < num_pstv; ++i)
                {
                    result.binnedDoseValues[i + 1] = doseVolumePairs[i].first;
                    result.accumulatedRelativeVolumes[i + 1] = doseVolumePairs[i].second;
                }

                break;  // skip all the rest
            }

            ///============================= 优化的粗略计算：CUDA版本
            /// 基于统计特征选择更合理的分箱数量; 提取剂量数据进行统计分析
            std::vector<double> doses(num_pstv);
            std::vector<double> volumes(num_pstv);
            double totalVolume = 0.0;

            for (size_t i = 0; i < num_pstv; ++i)
            {
                doses[i] = doseVolumePairs[i].first;
                volumes[i] = doseVolumePairs[i].second;
                totalVolume += volumes[i];
            }
#if 0
            size_t tmpi{};
            for (const auto d : doses)
            {
                if (fpclassify(d) == FP_ZERO)
                {
                    MYINFO("[my] id->{}, val->{}", tmpi, d);
                }
                tmpi++;
            }
#endif
            ///================================= CUDA内存分配

#if 10  /// padding
            double* d_doses;
            // 计算扩充后的大小，使其成为 warpSize(32)的整数倍
            const size_t padded_num_pstv = ((num_pstv + 31) / 32) * 32;
            CudaMaloc<double>(padded_num_pstv, d_doses);
            // 使用doses[0]填充扩充的部分
            if (padded_num_pstv > num_pstv && !doses.empty())
            {
                std::vector<double> pad_values(padded_num_pstv - num_pstv, doses[0]);
                CUDA_MCPYC(double, d_doses + num_pstv, pad_values.data(), pad_values.size(), cudaMemcpyHostToDevice);
            }
#else
            double* d_doses;
            CUDA_MCPYC(double, d_doses, doses.data(), num_pstv, cudaMemcpyHostToDevice);
#endif

            double *d_volumes, *d_minDose, *d_maxDose, *d_binVolumes;
            CudaMaloc<double>(num_pstv, d_volumes);
            CudaMaloc<double>(1, d_minDose, d_maxDose);
            /// 数据传输到设备
            CUDA_MCPYC(double, d_doses, doses.data(), num_pstv, cudaMemcpyHostToDevice);
            CUDA_MCPYC(double, d_volumes, volumes.data(), num_pstv, cudaMemcpyHostToDevice);
            CUDA_MCPYC(double, d_minDose, doses.data(), 1, cudaMemcpyHostToDevice);
            CUDA_MCPYC(double, d_maxDose, doses.data(), 1, cudaMemcpyHostToDevice);
            //MYINFO("[my] doses -> {}", doses);
            //MYINFO("[my] doses[0] -> {}", doses[0]);

            /// 计算剂量范围 - 使用CUDA
            // double dose = doses[idx];
            const int nBlock_F64_4 = (MYCEIL(MYCEIL(padded_num_pstv, BLOCK_SIZE), 4));  // 这里要除以4
            CuReduceMinMax<<<nBlock_F64_4, BLOCK_SIZE>>>(padded_num_pstv, d_doses, d_minDose, d_maxDose);
            CUDA_CHECK_LAST;

            /// 数据传输回主机
            double minDose, maxDose;
            CUDA_MCPYC(double, &minDose, d_minDose, 1, cudaMemcpyDeviceToHost);
            CUDA_MCPYC(double, &maxDose, d_maxDose, 1, cudaMemcpyDeviceToHost);
            //MYINFO("[my] minDose: {}, maxDose: {}, num_pstv: {}", minDose, maxDose, padded_num_pstv);
            minDose = std::max(0.0, minDose);

            /// 动态计算合适的分箱数
            const double doseRange = maxDose - minDose;

            ///============================= 粗略; 当剂量范围<=最小宽度, 报错
            if (doseRange <= doseMinWidth)
            {
                char errMsg[256];
                sprintf(errMsg,
                        "Error: dose range to small, range -> %f, min -> %f, dose -> %f; while voxels number too much -> (%zu); check input data.\n",
                        doseRange, minDose, maxDose, num_pstv);

                // 释放CUDA内存
                CudaFreeC(d_doses, d_volumes, d_minDose, d_maxDose);

                throw std::invalid_argument(errMsg);
            }

            ///============================= 当剂量范围>最小宽度，动态计算分箱数
            // 根据剂量范围和体素数量自适应确定分箱数
            int numBins = MIN_BINS;
            // 计算理论分箱数：每doseMinWidth剂量一个分箱; 加1.0 确保覆盖整个剂量范围
            const double theoreticalBins = doseRange / doseMinWidth + 1.0;
            // 将理论分箱数转换为整数，并限制在最小和最大分箱数之间
            const auto calculatedBins = static_cast<int>(theoreticalBins);
            numBins = std::max(MIN_BINS, std::min(MAX_BINS, calculatedBins));

            /// 准备分箱结果
            // 0 位置, n+1 位置, 两个哨兵
            constexpr size_t tickOfst = 1;
            result.binnedDoseValues.resize(numBins + 2, 0.0);
            result.accumulatedRelativeVolumes.resize(numBins + 2, 0.0);
            const double binWidth = doseRange / (numBins - 1);

            // 填充分箱剂量值
            for (int j = 0; j < numBins; ++j)
            {
                result.binnedDoseValues[j + tickOfst] = minDose + j * binWidth;
            }

            /// 为分箱体积分配CUDA内存
            CudaMaloc<double>(numBins + 2, d_binVolumes);
            CudaMSetC(0, numBins + 2, d_binVolumes);

            /// 统计每个分箱的体积 - 使用CUDA
            binVolumeKernel<<<nBlocks, BLOCK_SIZE>>>(d_doses, d_volumes, num_pstv, minDose, binWidth, numBins, d_binVolumes);
            CUDA_CHECK_LAST;

            /// 数据传输回主机
            CUDA_MCPYC(double, result.accumulatedRelativeVolumes.data(), d_binVolumes, (numBins + 2), cudaMemcpyDeviceToHost);

            /// 计算累积相对体积 - 使用CUDA
            CUDA_MCPYC(double, d_binVolumes, result.accumulatedRelativeVolumes.data(), (numBins + 2), cudaMemcpyHostToDevice);
            computeAccumulatedVolumesKernel<<<1, 1>>>(d_binVolumes, numBins);
            CUDA_CHECK_LAST;
            CUDA_MCPYC(double, result.accumulatedRelativeVolumes.data(), d_binVolumes, (numBins + 2), cudaMemcpyDeviceToHost);

            /// 哨兵 规整
            result.binnedDoseValues[0] = result.binnedDoseValues[1];
            result.binnedDoseValues.back() = result.binnedDoseValues[numBins];

            /// 释放CUDA内存
            CudaFreeC(d_doses, d_volumes, d_minDose, d_maxDose, d_binVolumes);

            break;

        } while (0);

#if 0
        for (const auto& v : result.accumulatedRelativeVolumes)
        {
            printf("v -> %f \n", v);
        }
#endif
        return result;
    }  // namespace dosestatcuda::DoseStatRet

    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  获取特定体积百分比对应的剂量 - 保持原有实现，因为这是一个查找操作，不适合并行化
    double GetMetric(double doseAtVolume, const std::vector<double>& binnedDoseValues, const std::vector<double>& accumulatedRelativeVolumes)
    {
        // 参数校验
        if (accumulatedRelativeVolumes.empty() || binnedDoseValues.empty())
        {
            MYDBG("Warning: Empty input vectors provided to GetMetric\n");
            return 0.0;
        }

        // 边界检查优化
        if (doseAtVolume >= accumulatedRelativeVolumes[0])
        {
            return binnedDoseValues[0];
        }
        if (doseAtVolume <= accumulatedRelativeVolumes.back())
        {
            return binnedDoseValues.back();
        }

        // 二分查找, 返回 所有等于 value 的元素范围
        auto [first, last] =
            std::equal_range(accumulatedRelativeVolumes.begin(), accumulatedRelativeVolumes.end(), doseAtVolume, std::greater<double>{});

        ///================= 存在恰好等于 dose 的情况 [v1,v2,...], rare
        if (first != last)
        {
            // 计算first到last范围内binnedDoseValues的平均值
            size_t count = 0;
            double sumDoses = 0.0;
            for (auto it = first; it != last; ++it)
            {
                const size_t index = it - accumulatedRelativeVolumes.begin();
                sumDoses += binnedDoseValues[index];
                count++;
            }
            const double avgDose = sumDoses / count;
            MYDBG("Volume -> %f, matching values -> (count: %zu), average dose -> %f \n", doseAtVolume, count, avgDose);
            return avgDose;
        }

        ///================= 不存在等于 dose 的情况, 插值; lower_bound == upper_bound
        const size_t ind = first - accumulatedRelativeVolumes.begin();
        MYDBG("Target dose -> %f, found index -> %zu\n", doseAtVolume, ind);
        ///------------- 边界检查
        if (ind == 0)
        {
            return binnedDoseValues[0];
        }
        if (ind >= accumulatedRelativeVolumes.size())
        {
            return binnedDoseValues.back();
        }
        size_t ind1 = ind - 1;  // little dose; big volume
        size_t ind2 = ind;      // big dose; small volume

        ///------------- 如果体积接近，返回平均剂量
        const double volDiff = accumulatedRelativeVolumes[ind2] - accumulatedRelativeVolumes[ind1];
        ///
        if (std::abs(volDiff) <= 1e-10)
        {
            MYDBG("Average; dose -> %f, index -> %zu\n", doseAtVolume, ind);
            return 0.5 * (binnedDoseValues[ind2] + binnedDoseValues[ind1]);
        }

        ///------------- 线性插值
        MYDBG("Linear interpolation; dose -> %f, index -> %zu\n", doseAtVolume, ind);
        double doseDiff = binnedDoseValues[ind2] - binnedDoseValues[ind1];
        double metric = binnedDoseValues[ind1] + (doseAtVolume - accumulatedRelativeVolumes[ind1]) * doseDiff / volDiff;
        return metric;
    }

    ///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  计算ROI内的剂量统计指标 - CUDA优化版本
    double GetDoseStatRoiInterp(const std::string& dose_type,
                                const std::vector<double>& binnedDoseValues,
                                const std::vector<double>& accumulatedRelativeVolumes,
                                const std::vector<double>& voxelVolume,
                                const std::vector<double>& voxelDose,
                                double doseAtVolume)
    {
        ///================= 计算平均剂量时，可以提前检查数组大小是否匹配
        if (voxelVolume.size() != voxelDose.size())
        {
            throw std::runtime_error("voxelVolume and voxelDose arrays have different sizes");
            return 0.0;
        }

        ///================= 使用哈希值或枚举类型替代字符串比较，提高性能
        if (dose_type == "average")
        {
            // CUDA优化的平均剂量计算
            size_t dataLen = voxelVolume.size();
            double totalVolume = 0.0;
            double weightedDose = 0.0;

            // 为大型数据集使用CUDA加速
            if (dataLen > 1000)
            {  // 阈值可以根据实际情况调整
                double *d_voxelVolume, *d_voxelDose, *d_totalVolume, *d_weightedDose;

                // 分配CUDA内存
                CudaMaloc<double>(dataLen, d_voxelVolume, d_voxelDose, d_totalVolume, d_weightedDose);

                // 初始化累加变量
                CudaMSetC(0, 1, d_totalVolume, d_weightedDose);

                // 数据传输到设备
                CUDA_MCPYC(double, d_voxelVolume, voxelVolume.data(), dataLen, cudaMemcpyHostToDevice);
                CUDA_MCPYC(double, d_voxelDose, voxelDose.data(), dataLen, cudaMemcpyHostToDevice);

                // 启动CUDA核函数
                int nBlocks = (dataLen + BLOCK_SIZE - 1) / BLOCK_SIZE;
                ComputeWeightedDoseKernel<<<nBlocks, BLOCK_SIZE>>>(d_voxelVolume, d_voxelDose, dataLen, d_totalVolume, d_weightedDose);
                CUDA_CHECK_LAST;

                // 数据传输回主机
                CUDA_MCPYC(double, &totalVolume, d_totalVolume, 1, cudaMemcpyDeviceToHost);
                CUDA_MCPYC(double, &weightedDose, d_weightedDose, 1, cudaMemcpyDeviceToHost);

                // 释放CUDA内存
                CudaFreeC(d_voxelVolume, d_voxelDose, d_totalVolume, d_weightedDose);
            }
            else
            {
                // 小型数据集使用CPU计算
                for (size_t i = 0; i < dataLen; ++i)
                {
                    totalVolume += voxelVolume[i];
                    weightedDose += voxelDose[i] * voxelVolume[i];
                }
            }

            if (totalVolume <= 0.0)
            {
                throw std::runtime_error("Prescription ROI volume is zero. Cannot scale to average dose.");
            }

            return weightedDose / totalVolume;
        }

        ///================= 提取目标体积百分比
        double targetDoseAtVolume;
        // 使用更高效的查找方式，例如unordered_map或者switch-case(如果使用枚举)
        if (dose_type == "D_98")
        {
            targetDoseAtVolume = 0.98;
        }
        else if (dose_type == "D_95")
        {
            targetDoseAtVolume = 0.95;
        }
        else if (dose_type == "D_50")
        {
            targetDoseAtVolume = 0.5;
        }
        else if (dose_type == "D_2")
        {
            targetDoseAtVolume = 0.02;
        }
        else if (dose_type == "DoseAtVolume")
        {
            targetDoseAtVolume = doseAtVolume;
        }
        else
        {
            throw std::runtime_error("Prescription type not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume' are supported");
        }

        // 调用优化后的get_metric函数
        return GetMetric(targetDoseAtVolume, binnedDoseValues, accumulatedRelativeVolumes);
    }
}  // namespace dosestatcuda
